#include "process_boxes_with_nms.h"
#include <algorithm>
#include <cmath>
#include <iostream>
#include <queue>

// float sigmoid(float x) { return 1.0f / (1.0f + std::exp(-x)); }

void softmax(const float *src, int len, float *dst) {
  float maxval = src[0];
  for (int i = 1; i < len; ++i)
    if (src[i] > maxval)
      maxval = src[i];
  float sum = 0.f;
  for (int i = 0; i < len; ++i) {
    dst[i] = std::exp(src[i] - maxval);
    sum += dst[i];
  }
  for (int i = 0; i < len; ++i)
    dst[i] /= sum;
}

float iou(const std::vector<float> &a, const std::vector<float> &b) {
  float x1 = std::max(a[0], b[0]);
  float y1 = std::max(a[1], b[1]);
  float x2 = std::min(a[2], b[2]);
  float y2 = std::min(a[3], b[3]);
  float iw = std::max(0.0f, x2 - x1);
  float ih = std::max(0.0f, y2 - y1);
  float inter = iw * ih;
  float area_a = (a[2] - a[0]) * (a[3] - a[1]);
  float area_b = (b[2] - b[0]) * (b[3] - b[1]);
  float union_ = area_a + area_b - inter;
  return union_ > 0 ? inter / union_ : 0.0f;
}

void nms_single_class(std::vector<Detection> &dets, float nms_thresh,
                      int max_det, std::vector<Detection> &result) {
  std::sort(
      dets.begin(), dets.end(),
      [](const Detection &a, const Detection &b) { return a.score > b.score; });
  std::vector<bool> removed(dets.size(), false);
  int count = 0;
  for (size_t i = 0; i < dets.size(); ++i) {
    if (removed[i])
      continue;
    result.push_back(dets[i]);
    if (++count >= max_det)
      break;
    for (size_t j = i + 1; j < dets.size(); ++j) {
      if (removed[j])
        continue;
      if (iou(dets[i].box, dets[j].box) > nms_thresh)
        removed[j] = true;
    }
  }
}

void process_boxes_with_nms(float *score[3], float *box[3], const int h[3],
                            const int w[3], float scale_x, float scale_y,
                            float conf_thresh, int reg_max, float nms_thresh,
                            int max_det, std::vector<Detection> &out_dets) {
  out_dets.clear();
  constexpr int n_classes = 6;
  // 1. 解码所有候选框
  std::vector<Detection> all_dets;

  for (int l = 0; l < 3; ++l) {
    int H = h[l], W = w[l];
    int stride = 8 << l;
    const float *score_ptr = score[l]; // [1, 6, H, W]
    const float *box_ptr = box[l];     // [1, 64, H, W]
    int area = H * W;
    std::vector<float> prob(reg_max);

    for (int y = 0; y < H; ++y) {
      for (int x = 0; x < W; ++x) {
        int offset = y * W + x;
        // 多目标解码，保留每个类别
        for (int c = 0; c < n_classes; ++c) {
          float cls_score = sigmoid(score_ptr[c * area + offset]);
          if (cls_score < conf_thresh)
            continue;

          // box_weights: [4]
          float box_weights[4] = {0};
          for (int b = 0; b < 4; ++b) {
            float logits[16];
            for (int k = 0; k < reg_max; ++k)
              logits[k] = box_ptr[(b * reg_max + k) * area + offset];
            softmax(logits, reg_max, prob.data());
            for (int k = 0; k < reg_max; ++k)
              box_weights[b] += prob[k] * k;
          }
          float x0 = box_weights[0];
          float y0 = box_weights[1];
          float x1 = box_weights[2];
          float y1 = box_weights[3];

          std::vector<float> box_vec = {(x + 0.5f - x0) * stride * scale_x,
                                        (y + 0.5f - y0) * stride * scale_y,
                                        (x + 0.5f + x1) * stride * scale_x,
                                        (y + 0.5f + y1) * stride * scale_y};
          all_dets.push_back({box_vec, cls_score, c});
        }
      }
    }
  }

  std::cout << "all_dets size: " << all_dets.size() << std::endl;
  // 2. 按类别分组做NMS
  for (int c = 0; c < 6; ++c) {
    std::vector<Detection> cls_dets;
    for (const auto &d : all_dets) {
      if (d.class_id == c)
        cls_dets.push_back(d);
    }
    if (!cls_dets.empty())
      nms_single_class(cls_dets, nms_thresh, max_det, out_dets);
  }
}